import os
import sys
sys.path.append("..")
import numpy as np
from transforms3d.euler import euler2mat,euler2quat,quat2euler


def lerp(euler0, euler1, alpha):
    """
    R0 * alpha + R1 * (1-alpha)
    """
    q0 = euler2quat(euler0[0]/180*np.pi, euler0[1]/180*np.pi, euler0[2]/180*np.pi, axes="sxyz")
    q1 = euler2quat(euler1[0]/180*np.pi, euler1[1]/180*np.pi, euler1[2]/180*np.pi, axes="sxyz")
    if np.dot(q0, -q1) > np.dot(q0, q1):
        q1 = -q1
    q = alpha * q0 + (1 - alpha) * q1
    q /= np.linalg.norm(q, ord=2)
    ai, aj, ak = quat2euler(q)
    return np.float32([ai/np.pi*180, aj/np.pi*180, ak/np.pi*180])


def VTSvisualization2pose(rigid2obj_path):
    T = np.eye(4)
    with open(rigid2obj_path, "r") as f:
        cnt = -1
        for line in f:
            cnt += 1
            vs = line.strip().split(",")
            values = np.float32([float(v) for v in vs])
            if cnt == 0:
                T[:3, 3] = values
            else:
                T[:3, :3] = euler2mat(values[0] / 180 * np.pi, values[1] / 180 * np.pi, values[2] / 180 * np.pi, axes="rxyz")

    return T


def fbx2proppose(fbx_path):  # 直接读出fbx里的pose数据
    flag = False  # flag=True则开始收集一批数据
    N = None
    cnt = -1
    
    # pose data
    timestamp = []
    rot = []
    trans = []

    with open(fbx_path) as f:
        for line in f:
            if ("KeyValueFloat:" in line) or (("KeyTime" in line) and len(timestamp) == 0):  # 一批数据的起始标记
                flag = True
                N_frame = int(line.strip().split(" ")[1][1:])
                if N is None:
                    N = N_frame
                else:
                    assert N == N_frame
                values = []
                N_values = 0
                continue
            if flag:
                line = line.strip()
                if " " in line:
                    line = line.split(" ")[-1]
                vs = line.split(",")
                for v in vs:
                    if v != "":
                        values.append(float(v))
                        N_values += 1
                if N_values == N:
                    # 得到了N帧的某维数据, 将其整合到pose中
                    cnt += 1
                    if (cnt == 0):
                        timestamp = values
                    elif cnt <= 3:
                        rot.append(np.float32(values))
                    else:
                        trans.append(np.float32(values) / 100)  # unit: cm -> m

                    # 从零开始收集下一批数据
                    flag = False
    
    rot = np.float32(rot).transpose(1, 0)
    trans = np.float32(trans).transpose(1, 0)
    
    # 检查时间戳是否连续
    td = timestamp[1] - timestamp[0]
    for i in range(N-1):
        if (timestamp[i+1]-timestamp[i] < td-100) or (td+100 < timestamp[i+1]-timestamp[i]):
            print("timestamp error: frame={}, relative_error={}".format(str(i), (timestamp[i+1]-timestamp[i])/(td+1e-6)))
    
    # 补齐数据
    rot0 = rot.copy()
    trans0 = trans.copy()
    rot, trans = [rot0[0]], [trans0[0]]
    N_total = N
    for i in range(1, N):
        if (timestamp[i]-timestamp[i-1] < td-100) or (td+100 < timestamp[i]-timestamp[i-1]):
            K = int((timestamp[i]-timestamp[i-1])/(td+1e-6) + 0.5)
            if i + 1 < N:
                target_rot = rot0[i+1]
                target_trans = trans0[i+1]
            else:
                target_rot = rot0[i]
                target_trans = trans0[i]
            N_total += K-1
            for k in range(K):
                rot.append(lerp(rot0[i], target_rot, 1-k/K))
                trans.append((1-k/K) * trans0[i] + (k/K) * target_trans)
        else:
            rot.append(rot0[i])
            trans.append(trans0[i])
    
    N = N_total
    rot = np.float32(rot)
    trans = np.float32(trans)
    assert (rot.shape == (N, 3)) and (trans.shape == (N, 3))

    proppose = np.zeros((N, 4, 4))  # obj2world
    for i in range(N):
        T = np.eye(4)  # rigid2world
        T[:3, :3] = euler2mat(rot[i, 0] / 180 * np.pi, rot[i, 1] / 180 * np.pi, rot[i, 2] / 180 * np.pi, axes="sxyz")
        T[:3, 3] = trans[i]
        proppose[i] = T
    
    return proppose


def fbx2objpose(fbx_path):  # 被交互的物体的pose, 由于VTS里导入fbx时把模型绕x轴旋转了, 所以这个函数里用inv_complement_T抵消这部分旋转
    """
    return: object pose for each frame, shape = (N, 4, 4), unit: m
    """

    inv_complement_T = np.eye(4)
    # inv_complement_T[:3, :3] = euler2mat(-np.pi / 2, 0, 0, axes="sxyz")

    proppose = fbx2proppose(fbx_path)
    assert proppose.shape[1:] == (4, 4)
    objpose = proppose @ inv_complement_T
    return objpose


def fbx2camerapose(fbx_path, camera2tracker_path):  # egocentric相机上的tracker的pose
    """
    return: egocentric camera pose for each frame, shape = (N, 4, 4), unit: m
    """

    camera2tracker = np.loadtxt(camera2tracker_path)

    proppose = fbx2proppose(fbx_path)
    assert proppose.shape[1:] == (4, 4)
    camerapose = proppose @ camera2tracker
    return camerapose


if __name__ == "__main__":

    objpose = fbx2objpose("../exp_data/20230323_data/20230323_debug2_xbox1.fbx")
    assert objpose.shape == (2001, 4, 4)
